{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "917d15c3-7eaf-423e-9860-c6bb5820aef3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, Dropout,Conv1D, AveragePooling1D, Flatten, Dense,Input\n",
    "from tensorflow.keras.optimizers import RMSprop\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.data import Dataset\n",
    "from tensorflow.keras.regularizers import l2\n",
    "import numpy as np\n",
    "import h5py\n",
    "from tqdm import trange\n",
    "import optuna\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import random\n",
    "import gc\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  # 0 表示第一个 GPU 设备，如果有多个 GPU，可以根据需要更改"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e90cd855-b90b-4a38-a0b0-046c212451b9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "def train_dataset(blockstart,blocknum,coeff_index):\n",
    "    X=[]\n",
    "    Y=[]\n",
    "\n",
    "    position=(np.load(r\"/mnt/riscv-trace/start_idx_10.npy\"))\n",
    "    # 切片位置,起点\n",
    "    for i in trange(blocknum,desc=\"Getting training data:\"):\n",
    "        data_infile=r\"/mnt/riscv-trace/metadata_part{}.npz\".format(blockstart+i)\n",
    "        # metadata url\n",
    "        data=np.load(data_infile,allow_pickle=True)\n",
    "        coefficient=data['allbytes']\n",
    "        trace_infile=r\"/mnt/riscv-trace/traces_part{}.npy\".format(blockstart+i)\n",
    "        # trace url\n",
    "        trace_file=np.load(trace_infile).astype(np.float32)\n",
    "\n",
    "        for t in range(trace_file.shape[0]):\n",
    "            for c in range(256):\n",
    "                r=c//4\n",
    " \n",
    "                for x in range(1):\n",
    "                    \n",
    "                    offset=0\n",
    "\n",
    "                    start=position[r]+offset\n",
    "                    end=start+1175\n",
    "\n",
    "                    t_slice=trace_file[t][start:end]\n",
    "\n",
    "                    if coeff_index==0 and c%4==0:\n",
    "                        if coefficient[t][c*8:c*8+8][:2]==\"00\":\n",
    "                            Y.append(0)\n",
    "                            X.append(t_slice)\n",
    "                        else:\n",
    "                            Y.append(1)\n",
    "                            X.append(t_slice)\n",
    "            \n",
    "                    if coeff_index==1 and c%4==1:\n",
    "                        if coefficient[t][c*8:c*8+8][:2]==\"00\":\n",
    "                            Y.append(0)\n",
    "                            X.append(t_slice)\n",
    "                        else:\n",
    "                            Y.append(1)\n",
    "                            X.append(t_slice)\n",
    "                            \n",
    "                    if coeff_index==2 and c%4==2:\n",
    "                        if coefficient[t][c*8:c*8+8][:2]==\"00\":\n",
    "                            Y.append(0)\n",
    "                            X.append(t_slice)\n",
    "                        else:\n",
    "                            Y.append(1)\n",
    "                            X.append(t_slice)\n",
    "                            \n",
    "                    if coeff_index==3 and c%4==3:\n",
    "                        if coefficient[t][c*8:c*8+8][:2]==\"00\":\n",
    "                            Y.append(0)\n",
    "                            X.append(t_slice)\n",
    "                        else:\n",
    "                            Y.append(1)\n",
    "                            X.append(t_slice)\n",
    "            \n",
    "\n",
    "    new_x = np.array(X)\n",
    "    new_y = np.array(Y)\n",
    "    y_train = tf.keras.utils.to_categorical(new_y, num_classes=2, dtype=int)          \n",
    "    return new_x,y_train           \n",
    "                     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9955d26-5c49-42c7-a9f1-db366d88e905",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    " # Trial 1 finished with value: 0.9970999956130981 and parameters: \n",
    " #        {'filters_1': 248, 'kernel_size_1': 9, 'pool_size_1': 5, \n",
    " #         'filters_2': 203, 'kernel_size_2': 13, 'pool_size_2': 3, \n",
    " #         'filters_3': 115, 'kernel_size_3': 7, 'pool_size_3': 5, \n",
    " #         'dense_units1': 218, 'dropout_rate1': 0.15879146923376725, 'learning_rate': 0.0004239454853083666}. \n",
    " #        Best is trial 1 with value: 0.9970999956130981.\n",
    "\n",
    "def riscv_model(input_dim):\n",
    "    model = Sequential()\n",
    "\n",
    "    # Block 1\n",
    "    model.add(Conv1D(filters=248, kernel_size=9, activation='relu', padding='same', name='block1_conv1', input_shape=(input_dim, 1)))\n",
    "    model.add(AveragePooling1D(pool_size=5, strides=2, name='block1_pool'))\n",
    "\n",
    "    # Block 2\n",
    "    model.add(Conv1D(filters=203, kernel_size=13, activation='relu', padding='same', name='block2_conv1'))\n",
    "    model.add(AveragePooling1D(pool_size=3, strides=2, name='block2_pool'))\n",
    "    \n",
    "    # Block 3\n",
    "    model.add(Conv1D(filters=115, kernel_size=7, activation='relu', padding='same', name='block3_conv1'))\n",
    "    model.add(AveragePooling1D(pool_size=5, strides=2, name='block3_pool'))\n",
    "\n",
    "\n",
    "    # Flatten and Dense layers\n",
    "    model.add(Flatten(name='flatten'))\n",
    "    \n",
    "    model.add(Dense(218, activation='relu'))\n",
    "    model.add(Dropout(0.15879146923376725))\n",
    "\n",
    "    # model.add(Dense(dense_units2, activation='relu'))\n",
    "    # model.add(Dropout(dropout_rate2))\n",
    "    \n",
    "    model.add(Dense(2, activation='softmax'))\n",
    "\n",
    "    optimizer = RMSprop(0.0004239454853083666)\n",
    "    model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])\n",
    "    return model\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ed89124e-d6eb-40e0-86e4-21ef192fbb53",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting training data:: 100%|██████████| 25/25 [00:35<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 512000 samples, validate on 128000 samples\n",
      "Epoch 1/2\n",
      "512000/512000 [==============================] - 143s 280us/sample - loss: 0.0477 - acc: 0.9931 - val_loss: 0.0130 - val_acc: 0.9963\n",
      "Epoch 2/2\n",
      "512000/512000 [==============================] - 148s 290us/sample - loss: 0.0117 - acc: 0.9965 - val_loss: 0.0083 - val_acc: 0.9970\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting training data:: 100%|██████████| 25/25 [00:39<00:00,  1.57s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 512000 samples, validate on 128000 samples\n",
      "Epoch 1/2\n",
      "512000/512000 [==============================] - 142s 278us/sample - loss: 0.0702 - acc: 0.9895 - val_loss: 0.0143 - val_acc: 0.9958\n",
      "Epoch 2/2\n",
      "512000/512000 [==============================] - 149s 290us/sample - loss: 0.0111 - acc: 0.9959 - val_loss: 0.0093 - val_acc: 0.9963\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting training data:: 100%|██████████| 25/25 [00:35<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 512000 samples, validate on 128000 samples\n",
      "Epoch 1/2\n",
      "512000/512000 [==============================] - 144s 281us/sample - loss: 0.0404 - acc: 0.9926 - val_loss: 0.0112 - val_acc: 0.9960\n",
      "Epoch 2/2\n",
      "512000/512000 [==============================] - 149s 290us/sample - loss: 0.0106 - acc: 0.9966 - val_loss: 0.0071 - val_acc: 0.9974\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting training data:: 100%|██████████| 25/25 [00:35<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 512000 samples, validate on 128000 samples\n",
      "Epoch 1/2\n",
      "512000/512000 [==============================] - 144s 281us/sample - loss: 0.0273 - acc: 0.9947 - val_loss: 0.0108 - val_acc: 0.9970\n",
      "Epoch 2/2\n",
      "512000/512000 [==============================] - 149s 291us/sample - loss: 0.0079 - acc: 0.9974 - val_loss: 0.0122 - val_acc: 0.9970\n"
     ]
    }
   ],
   "source": [
    "'''train'''\n",
    "for i in range(4):\n",
    "    coeff_index=i\n",
    "    model=riscv_model(1175)\n",
    "    # model=simple_mlp_model(760,2)\n",
    "    X,Y=train_dataset(blockstart=0,blocknum=25,coeff_index=coeff_index)\n",
    "    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)\n",
    "\n",
    "\n",
    "    X_train=X_train.reshape([X_train.shape[0],X_train.shape[1],1])\n",
    "    # class_weight={0:2.4874611344584996 , 1:1},\n",
    "    model.fit(X_train, Y_train, epochs=2,  batch_size=512,validation_split=0.2)\n",
    "    model.save(r\"/mnt/riscv-models/model_coeff{}.h5\".format(coeff_index), save_format='h5')\n",
    "    del X\n",
    "    del Y\n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "da4029b3-7a6c-4d6e-98c0-49c585939ddb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Trial 0 finished with value: 0.9968888759613037 and parameters: {'filters_1': 19, 'kernel_size_1': 18, 'pool_size_1': 3, 'filters_2': 24, 'kernel_size_2': 15, 'pool_size_2': 3, \n",
    "# 'filters_3': 102, 'kernel_size_3': 7, 'pool_size_3': 3, 'dense_units1': 44, \n",
    "# 'dropout_rate1': 0.498730101923068, 'weight': 4.7693248391574725, 'learning_rate': 0.0004075973297959411}. Best is trial 0 with value: 0.9968888759613037.\n",
    "def objective(trial):\n",
    "    \n",
    "    input_dim=1175\n",
    "    \n",
    "    filters_1 = trial.suggest_int('filters_1', 16, 256)\n",
    "    kernel_size_1 = trial.suggest_int('kernel_size_1', 2, 15)\n",
    "    pool_size_1 = trial.suggest_int('pool_size_1', 2, 5)\n",
    "\n",
    "    filters_2 = trial.suggest_int('filters_2', 16, 256)\n",
    "    kernel_size_2 = trial.suggest_int('kernel_size_2', 2, 15)\n",
    "    pool_size_2 = trial.suggest_int('pool_size_2', 2, 5)\n",
    "\n",
    "    filters_3 = trial.suggest_int('filters_3', 16, 256)\n",
    "    kernel_size_3 = trial.suggest_int('kernel_size_3', 2, 15)\n",
    "    pool_size_3 = trial.suggest_int('pool_size_3', 2, 5)\n",
    "    \n",
    "    \n",
    "    dense_units1 = trial.suggest_int('dense_units1', 10, 500)\n",
    "    dropout_rate1 = trial.suggest_uniform('dropout_rate1', 0.1, 0.4)\n",
    "\n",
    "    \n",
    "    # Additional hyperparameters\n",
    "    # weight = trial.suggest_loguniform('weight', 1, 6)\n",
    "    learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 1e-2)\n",
    "    \n",
    "    model = Sequential()\n",
    "\n",
    "    # Block 1\n",
    "    model.add(Conv1D(filters=filters_1, kernel_size=kernel_size_1, activation='relu', padding='same', name='block1_conv1', input_shape=(input_dim, 1)))\n",
    "    model.add(AveragePooling1D(pool_size=pool_size_1, strides=2, name='block1_pool'))\n",
    "\n",
    "    # Block 2\n",
    "    model.add(Conv1D(filters=filters_2, kernel_size=kernel_size_2, activation='relu', padding='same', name='block2_conv1'))\n",
    "    model.add(AveragePooling1D(pool_size=pool_size_2, strides=2, name='block2_pool'))\n",
    "\n",
    "    # Block 3\n",
    "    model.add(Conv1D(filters=filters_3, kernel_size=kernel_size_3, activation='relu', padding='same', name='block3_conv1'))\n",
    "    model.add(AveragePooling1D(pool_size=pool_size_3, strides=2, name='block3_pool'))\n",
    "    \n",
    "    # Flatten and Dense layers\n",
    "    model.add(Flatten(name='flatten'))\n",
    "    \n",
    "    model.add(Dense(dense_units1, activation='relu'))\n",
    "    model.add(Dropout(dropout_rate1))\n",
    "\n",
    "    \n",
    "    model.add(Dense(2, activation='softmax'))\n",
    "    \n",
    "    optimizer = RMSprop(learning_rate)\n",
    "    model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])\n",
    "    \n",
    "\n",
    "    X,Y=train_dataset(blockstart=0,blocknum=25,coeff_index=0)\n",
    "    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42)\n",
    "\n",
    "\n",
    "    X_train=X_train.reshape([X_train.shape[0],input_dim,1])\n",
    "    model.fit(X_train, Y_train, epochs=2, batch_size=500, validation_split=0.25, verbose=1)\n",
    "\n",
    "    # Evaluate the model\n",
    "    X_test=X_test.reshape([X_test.shape[0],input_dim,1])\n",
    "    accuracy = model.evaluate(X_test, Y_test)[1]\n",
    "    del X\n",
    "    del Y\n",
    "    gc.collect()\n",
    "    return accuracy\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c10956d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()\n",
    "\n",
    "\n",
    "study = optuna.create_study(direction='maximize')\n",
    "study.optimize(objective, n_trials=20)\n",
    "\n",
    "# Print the best hyperparameters\n",
    "best_params = study.best_params\n",
    "print(\"Best Hyperparameters:\", best_params)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myconda",
   "language": "python",
   "name": "myconda"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
